from myapp.pp_py.src.data_structures.attribute_structures import *
from myapp.pp_py.src.data_structures.base_structures import Map, HexCell
import shapefile
from tqdm import tqdm
import os
import sys
import pickle

from myapp.pp_py.src.data_structures.weight_structures import TimeWeight, ConsumptionWeight
from myapp.pp_py.src.enums.pp_enum import AttributeType, RoadType


def map_to_shp(map: Map, path: str):
    """将Map写入Shapefile文件中"""

    # 创建 shapefile writer
    with shapefile.Writer(path) as shp:
        # 定义字段
        shp.field("index", "C", size=20)  # 格网索引
        shp.field("elevation", "F", decimal=4)  # 高程
        shp.field("slope", "F", decimal=4)  # 坡度
        shp.field("cv", "F", decimal=4)  # 高程变异系数
        shp.field("roughness", "F", decimal=4)  # 粗糙度
        shp.field("relief", "F", decimal=4)  # 地形起伏
        shp.field("curvature", "F", decimal=4)  # 曲率
        shp.field("exposure", "F", decimal=4)  # 坡向
        shp.field("terrain_type", "N", decimal=0)  # 地形类型(水体、建筑物等)
        shp.field("geology", "C", size=20)  # 地质成分
        shp.field("time_factor", "F", decimal=4)  # 时间因子
        shp.field("consumption_factor", "F", decimal=4)  # 能耗因子
        shp.field("rejected", "F", decimal=1)  # 是否拒绝

        def get_terrain_type(cell: HexCell):
            """
            获取地形类型
            :param cell: Cell 对象
            :return: 地形类型的整数值，如果不存在则返回 0
            """

            for attr in cell.attributes:
                if isinstance(attr, Water):
                    return attr.get_sub_attribute(Water.Type).value if attr.get_sub_attribute(
                        Water.Type) else AttributeType.UNKNOWN.value
                elif isinstance(attr, Building):
                    return attr.get_sub_attribute(Building.Type).value if attr.get_sub_attribute(
                        Building.Type) else AttributeType.UNKNOWN.value
                elif isinstance(attr, Vegetation):
                    return attr.get_sub_attribute(Vegetation.Type).value if attr.get_sub_attribute(
                        Vegetation.Type) else AttributeType.UNKNOWN.value
                elif isinstance(attr, Soil):
                    return attr.get_sub_attribute(Soil.Type).value if attr.get_sub_attribute(
                        Soil.Type) else AttributeType.UNKNOWN.value

            # 如果没有找到地形类型，则检查道路类型
            if cell.road_type != RoadType.NOWAY.value:
                return AttributeType.ROAD.value
            return None  # 默认返回 None，如果没有找到地形类型

        def get_geology_type(cell: HexCell):
            for attr in cell.attributes:
                if isinstance(attr, Geology):
                    return attr.get_sub_attribute(Geology.Lithology).value if attr.get_sub_attribute(
                        Geology.Lithology) else AttributeType.UNKNOWN.value
            return None  # 默认返回 None，如果没有找到地质成分

        # 遍历 Map 中的 Cell 对象
        for cell in tqdm(map.cells.values(), desc="将map转为shp:"):
            cell: HexCell
            # 写入属性
            record = [
                cell.index,  # h3_index
                cell.get_attribute(DemAttribute.Elevation).get_value() if cell.get_attribute(
                    DemAttribute.Elevation) else None,
                cell.get_attribute(DemAttribute.Slope).get_value() if cell.get_attribute(DemAttribute.Slope) else None,
                cell.get_attribute(DemAttribute.ElevationCoefficientOfVariation).get_value() if cell.get_attribute(
                    DemAttribute.ElevationCoefficientOfVariation) else None,
                cell.get_attribute(DemAttribute.Roughness).get_value() if cell.get_attribute(
                    DemAttribute.Roughness) else None,
                cell.get_attribute(DemAttribute.Relief).get_value() if cell.get_attribute(
                    DemAttribute.Relief) else None,
                cell.get_attribute(DemAttribute.Curvature).get_value() if cell.get_attribute(
                    DemAttribute.Curvature) else None,
                cell.get_attribute(DemAttribute.Exposure).get_value() if cell.get_attribute(
                    DemAttribute.Exposure) else None,
                get_terrain_type(cell),  # 地形类型
                get_geology_type(cell),  # 地质成分
                cell.get_weight(TimeWeight).get_value() if cell.get_weight(TimeWeight) else None,  # 时间因子
                cell.get_weight(ConsumptionWeight).get_value() if cell.get_weight(ConsumptionWeight) else None,  # 能耗因子
                1 if getattr(cell, "rejected", False) else 0  # 是否拒绝
            ]
            shp.record(*record)
            # 写入几何形状
            shp.poly([[(lon, lat) for lat, lon in cell.vertices]])

    # 为 Shapefile 文件生成 WGS84 坐标系的 .prj 文件
    wgs84_prj = """GEOGCS["WGS 84",
        DATUM["WGS_1984",
            SPHEROID["WGS 84",6378137,298.257223563]],
        PRIMEM["Greenwich",0],
        UNIT["degree",0.0174532925199433]]"""
    prj_path = os.path.splitext(path)[0] + ".prj"
    with open(prj_path, "w") as prj_file:
        prj_file.write(wgs84_prj)


def save_map(map: Map, path: str):
    """使用pickle保存map为bin文件"""
    with open(path, 'wb') as f:
        pickle.dump(map, f)


def load_map(bin_path='data/玄武区.bin') -> Map:
    """使用pickle加载map"""
    # 处理模块重命名问题
    sys.modules['src.data_structures'] = sys.modules['myapp.pp_py.src.data_structures']
    sys.modules['src.data_structures.base_structures'] = sys.modules['myapp.pp_py.src.data_structures.base_structures']
    sys.modules['src.data_structures.attribute_structures'] = sys.modules[
        'myapp.pp_py.src.data_structures.attribute_structures']
    sys.modules['src.data_structures.weight_structures'] = sys.modules[
        'myapp.pp_py.src.data_structures.weight_structures']
    sys.modules['src.enums.pp_enum'] = sys.modules['myapp.pp_py.src.enums.pp_enum']

    with open(bin_path, 'rb') as f:
        map = pickle.load(f)
    return map